# -*- coding: utf-8 -*-
"""
@Time    : 2024/10/22 11:25 
@Author  : ZhangShenao 
@File    : retriever.py 
@Desc    : 检索器

执行Hybrid Search混合检索(多路召回)+Reranking重排序
采用VectorSearch向量检索+BM25关键词检索策略
可以兼具关键词检索的准确率和向量检索的上下文相关性
并基于Reranking重排序策略,将最相关、最重要的文档置于前列
"""
from FlagEmbedding import FlagReranker
from langchain.retrievers import EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStore


def retrieving(query: str,
               docs: list[Document],
               vector_store: VectorStore,
               reranker: FlagReranker,
               top_k: int = 5) -> list[Document]:
    """
    Retrieving 检索过程
    :param query: 原始提问
    :param docs: 文档语料列表
    :param vector_store: 向量数据库
    :param reranker: 重排序器
    :param top_k: 检索最相近的top_k个文档,默认为5
    :return: 检索结果
    """

    print("Retrieving开始执行!")
    print(f"原始提问: {query} , 检索最相关的 {top_k} 个结果")

    # 执行Hybrid Search混合检索
    hybrid_search_docs = hybrid_search(query=query, docs=docs, vector_store=vector_store, top_k=top_k)
    print(f"Hybrid Search文档数量: {len(hybrid_search_docs)}")

    # 执行Reranking重排序
    reranking_docs = rerank(query=query, docs=hybrid_search_docs, top_k=top_k, reranker=reranker)
    print(f"Reranking文档数量: {len(reranking_docs)}")

    # 返回重排序后的结果
    print("Retrieving执行完成!")
    return reranking_docs


def hybrid_search(query: str, docs: list[Document], vector_store: VectorStore,
                  top_k: int = 5) -> list[Document]:
    """
    混合检索(多路召回) 采用VectorSearch向量检索+BM25关键词检索策略
    :param query: 原始提问
    :param docs: 文档语料列表
    :param vector_store: 向量数据库
    :param top_k: 检索最相近的top_k个文档,默认为5
    :return: 检索结果
    """

    print("Hybrid Search开始执行!")
    print(f"原始提问: {query} , 检索最相关的 {top_k} 个结果")

    # 创建BM25关键词检索器,并指定top_k
    bm25_retriever = BM25Retriever.from_documents(documents=docs)
    bm25_retriever.k = top_k

    # 创建VectorStore向量检索器,并指定相似度得分阈值和top_k
    vector_store_retriever = vector_store.as_retriever(
        search_type="similarity_score_threshold",
        search_kwargs={
            "k": top_k,
            "score_threshold": 0.5,
        }
    )

    # 创建EnsembleRetriever集成检索器(混合检索器)
    # 结合关键词检索与向量检索,实现多路召回策略
    # 可以为每个检索器指定权重
    hybrid_retriever = EnsembleRetriever(
        retrievers=[bm25_retriever, vector_store_retriever],
        weights=[0.5, 0.5],
    )

    # 执行混合检索
    result = hybrid_retriever.invoke(query)

    print("Hybrid Search执行完成!")
    return result


def rerank(query: str, docs: list[Document], top_k: int, reranker: FlagReranker) -> list[Document]:
    """
    Reranking重排序策略: 利用语义相似性计算,将最相关、最重要的文档置于前列
    :param query: 原始提问
    :param docs: 文档语料列表
    :param top_k: 检索最相近的top_k个文档
    :param reranker: 重排序器
    :return: 重排序后的结果
    """

    print("Rerank开始执行!")

    # 构造输入对,原始query与每个文档都构成一个输入对
    input_pairs = [(query, doc) for doc in docs]

    # 计算每个doc与query之间的相似性得分
    scores = reranker.compute_score(input_pairs, normalize=True)
    print(f"Rerank得分: {scores}")

    # 对得分进行排序,并获取排名前top_k的docs
    sorted_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
    reranking_docs = [docs[i] for i in sorted_indices[:top_k]]

    # 返回重排序后的结果
    print("Rerank执行完成!")
    return reranking_docs
